[Feature] Qwen3 next support sp. #3225
Conversation
Signed-off-by: weijinqian_v1 <[email protected]>
Signed-off-by: weijinqian_v1 <[email protected]>
Signed-off-by: weijinqian_v1 <[email protected]>
Signed-off-by: weijinqian_v1 <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for sequence parallelism (SP) for the Qwen3-Next model. The changes involve refactoring linear parallel operations, adding SP-specific logic in the model's decoder layer, and introducing a new AscendGemmaRMSNorm layer. My review has identified two critical bugs: one in the new AscendGemmaRMSNorm layer definition that would cause runtime errors due to an incorrect method signature, and another in the refactored linear op dispatcher which has a function call signature mismatch, also leading to a TypeError. These issues need to be addressed to ensure correctness.
| @staticmethod | ||
| def forward_oot( | ||
| self, | ||
| variance_epsilon: float, | ||
| x: torch.Tensor, | ||
| residual: Optional[torch.Tensor] = None, | ||
| ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: | ||
| """PyTorch-native implementation equivalent to forward().""" | ||
| return self.forward_static(self.weight.data, self.variance_epsilon, x, | ||
| residual) |
There was a problem hiding this comment.
The forward_oot method in AscendGemmaRMSNorm is incorrectly defined as a static method with a mismatched signature. It should be an instance method, and its signature should match the forward method it's overriding, which is (self, x, residual=None). The current implementation will lead to incorrect argument passing at runtime, as the variance_epsilon parameter in the signature will receive the x tensor, and the x parameter will receive the residual tensor. The variance_epsilon should be accessed from self.variance_epsilon instead of being a parameter.
| @staticmethod | |
| def forward_oot( | |
| self, | |
| variance_epsilon: float, | |
| x: torch.Tensor, | |
| residual: Optional[torch.Tensor] = None, | |
| ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: | |
| """PyTorch-native implementation equivalent to forward().""" | |
| return self.forward_static(self.weight.data, self.variance_epsilon, x, | |
| residual) | |
| def forward_oot( | |
| self, | |
| x: torch.Tensor, | |
| residual: Optional[torch.Tensor] = None, | |
| ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: | |
| """PyTorch-native implementation equivalent to forward().""" | |
| return self.forward_static(self.weight.data, self.variance_epsilon, x, | |
| residual) |
vllm_ascend/ops/linear_op.py
Outdated
| def _get_column_parallel_op( | ||
| disable_tp, prefix, layer | ||
| ) -> Tuple[Optional[Union[MLPColumnParallelOp, SequenceMergedColumnParallelOp, | ||
| SequenceQKVParallelOp]], int, int]: |
There was a problem hiding this comment.
The signature of _get_column_parallel_op is (disable_tp, prefix, layer), but it is called as _get_column_parallel_op(prefix, layer) in get_parallel_op at line 463. This mismatch will cause a TypeError at runtime. The disable_tp parameter is not used within the function and should be removed from the signature to match the call site.
| def _get_column_parallel_op( | |
| disable_tp, prefix, layer | |
| ) -> Tuple[Optional[Union[MLPColumnParallelOp, SequenceMergedColumnParallelOp, | |
| SequenceQKVParallelOp]], int, int]: | |
| def _get_column_parallel_op( | |
| prefix, layer | |
| ) -> Tuple[Optional[Union[MLPColumnParallelOp, SequenceMergedColumnParallelOp, | |
| SequenceQKVParallelOp]], int, int]: |
Signed-off-by: weijinqian_v1 <[email protected]>
Signed-off-by: weijinqian_v1 <[email protected]>
Signed-off-by: weijinqian_v1 <[email protected]>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: weijinqian_v1 <[email protected]>
Uh oh!
There was an error while loading. Please reload this page.